feat(python): payload codecs with global chain and per-task registry#365
Conversation
|
Warning Review limit reachedYou’ve reached a temporary PR review limit under our Fair Usage Limits Policy. Next review available in: 30 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughIntroduces a composable payload codec framework (Gzip, AES-GCM, HMAC codecs plus a CodecSerializer wrapper and CryptoError) integrated into Queue's construction, enqueue/dedup ordering, per-task codec registration, workflow/fan-out payload encoding, and result serialization across native async, async executor, and prefork worker paths, with accompanying tests. ChangesPayload codec framework and queue integration
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant Queue
participant CodecChain
participant Storage
Caller->>Queue: enqueue(task_name, args, kwargs)
Queue->>Queue: serialize (args, kwargs)
Queue->>Queue: compute dedup key from pre-codec bytes
Queue->>CodecChain: _apply_task_codecs(task_name, payload)
CodecChain-->>Queue: encoded payload
Queue->>Queue: validate payload size (post-codec)
Queue->>Storage: persist job with encoded payload
sequenceDiagram
participant Executor
participant QueueRef
participant Cloudpickle
Executor->>QueueRef: check _queue_ref
alt queue present
Executor->>QueueRef: _serialize_result(task_name, result)
QueueRef-->>Executor: serialized bytes
else no queue
Executor->>Cloudpickle: dumps(result)
Cloudpickle-->>Executor: serialized bytes
end
Possibly related PRs
Suggested labels: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
sdks/python/taskito/app.py (1)
343-365: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
_encode_payload()should enforce the payload cap itselfSeveral callers bypass the external
_check_payload_size()step (sdks/python/taskito/mixins/predicates.py::_reenqueue_after_defer,sdks/python/taskito/workflows/builder.py::_compile, andsdks/python/taskito/workflows/tracker/fan_out.py), so oversized task payloads can still reach the backend. Moving the size check into_encode_payload()would make the limit consistent and let_dispatch_batched_payload()drop its extra guard.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@sdks/python/taskito/app.py` around lines 343 - 365, `_encode_payload()` in `app.py` should enforce `max_payload_bytes` directly so oversized payloads are rejected even when callers skip `_check_payload_size()`. Add the size check after serialization/codecs in `_encode_payload()` using the existing `_check_payload_size()` helper and task name, and then remove the redundant guard from `_dispatch_batched_payload()` so all paths go through the same limit enforcement.
🧹 Nitpick comments (3)
crates/taskito-python/src/native_async/task_executor.rs (1)
152-165: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated result-serialization block with
py_worker.rs.This exact block (None-check,
_queue_reflookup,_serialize_result/cloudpickle.dumpsfallback,.extract()) is duplicated verbatim incrates/taskito-python/src/py_worker.rs(lines 210-223). Consider extracting a shared helper (e.g.serialize_task_result(context_mod, job, result) -> PyResult<Option<Vec<u8>>>) in a common module so both worker paths stay in sync as this codec-aware logic evolves.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/taskito-python/src/native_async/task_executor.rs` around lines 152 - 165, The result-serialization logic in task_executor.rs is duplicated in py_worker.rs, so extract the None-check, _queue_ref lookup, _serialize_result/cloudpickle.dumps fallback, and byte extraction into a shared helper such as serialize_task_result(context_mod, job, result) returning PyResult<Option<Vec<u8>>>. Update both TaskExecutor and the corresponding py_worker path to call the shared helper so the codec-aware behavior stays consistent in one place.crates/taskito-python/src/py_worker.rs (1)
210-223: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSame duplicated serialization block as
native_async/task_executor.rs.See the companion comment on
crates/taskito-python/src/native_async/task_executor.rs(lines 152-165) — this block is byte-for-byte identical. Worth extracting to a shared helper to avoid drift between the two worker paths.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/taskito-python/src/py_worker.rs` around lines 210 - 223, The result serialization logic in `py_worker.rs` is duplicated verbatim in the native async worker path, so extract the shared behavior into a common helper and have both `py_worker` and `native_async/task_executor` call it. Keep the helper responsible for the `result.is_none()` handling, `_queue_ref` lookup, `_serialize_result` vs `cloudpickle.dumps` fallback, and `bytes` extraction so the two worker paths stay in sync.sdks/python/taskito/mixins/decorators.py (1)
641-711: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win
periodic()doesn't expose acodecs=param, so scheduled tasks can't opt into per-task codecs.Both the plain-task branch (line 694:
self.task(name=name, queue=queue)(fn)) and the workflow-launcher branch callself.task(...)without forwarding acodecs=list, so a periodic task can never get an entry in_task_codecs, even though_encode_payloadis now used to build its payload. This mirrors the pre-existingserializer=gap, but is worth closing now that codec support is being wired through the queue.♻️ Possible fix
def periodic( self, cron: str, name: str | None = None, args: tuple = (), kwargs: dict | None = None, queue: str = "default", timezone: str | None = None, + codecs: list[str] | None = None, ) -> Callable[[Callable[..., Any]], TaskWrapper]: ... - wrapper = self.task(name=name, queue=queue)(fn) + wrapper = self.task(name=name, queue=queue, codecs=codecs)(fn)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@sdks/python/taskito/mixins/decorators.py` around lines 641 - 711, The periodic decorator in `periodic()` is not forwarding per-task codec configuration, so scheduled tasks can’t register their codecs even though payloads are encoded through `_encode_payload`. Update both branches inside `decorator()`—the normal task registration via `self.task(...)` and the workflow launcher registration via `@self.task(...)`—to accept and pass through a `codecs=` parameter. Make sure the codec list is propagated into the underlying task registration so entries are created alongside `_periodic_configs` and the generated launcher task stays consistent with `TaskWrapper` setup.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@sdks/python/taskito/app.py`:
- Around line 343-365: `_encode_payload()` in `app.py` should enforce
`max_payload_bytes` directly so oversized payloads are rejected even when
callers skip `_check_payload_size()`. Add the size check after
serialization/codecs in `_encode_payload()` using the existing
`_check_payload_size()` helper and task name, and then remove the redundant
guard from `_dispatch_batched_payload()` so all paths go through the same limit
enforcement.
---
Nitpick comments:
In `@crates/taskito-python/src/native_async/task_executor.rs`:
- Around line 152-165: The result-serialization logic in task_executor.rs is
duplicated in py_worker.rs, so extract the None-check, _queue_ref lookup,
_serialize_result/cloudpickle.dumps fallback, and byte extraction into a shared
helper such as serialize_task_result(context_mod, job, result) returning
PyResult<Option<Vec<u8>>>. Update both TaskExecutor and the corresponding
py_worker path to call the shared helper so the codec-aware behavior stays
consistent in one place.
In `@crates/taskito-python/src/py_worker.rs`:
- Around line 210-223: The result serialization logic in `py_worker.rs` is
duplicated verbatim in the native async worker path, so extract the shared
behavior into a common helper and have both `py_worker` and
`native_async/task_executor` call it. Keep the helper responsible for the
`result.is_none()` handling, `_queue_ref` lookup, `_serialize_result` vs
`cloudpickle.dumps` fallback, and `bytes` extraction so the two worker paths
stay in sync.
In `@sdks/python/taskito/mixins/decorators.py`:
- Around line 641-711: The periodic decorator in `periodic()` is not forwarding
per-task codec configuration, so scheduled tasks can’t register their codecs
even though payloads are encoded through `_encode_payload`. Update both branches
inside `decorator()`—the normal task registration via `self.task(...)` and the
workflow launcher registration via `@self.task(...)`—to accept and pass through
a `codecs=` parameter. Make sure the codec list is propagated into the
underlying task registration so entries are created alongside
`_periodic_configs` and the generated launcher task stays consistent with
`TaskWrapper` setup.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: ad0e23e8-a503-4931-88dc-f84c17de9817
📒 Files selected for processing (16)
crates/taskito-python/src/native_async/task_executor.rscrates/taskito-python/src/py_worker.rssdks/python/taskito/__init__.pysdks/python/taskito/app.pysdks/python/taskito/async_support/executor.pysdks/python/taskito/codecs.pysdks/python/taskito/exceptions.pysdks/python/taskito/mixins/decorators.pysdks/python/taskito/mixins/predicates.pysdks/python/taskito/prefork/child.pysdks/python/taskito/workflows/builder.pysdks/python/taskito/workflows/saga/orchestrator.pysdks/python/taskito/workflows/tracker/fan_out.pysdks/python/tests/core/test_codecs.pysdks/python/tests/core/test_result_serialization.pysdks/python/tests/worker/test_native_async.py
Summary
Adds a payload-codec layer to the Python SDK with wire formats that are part of the cross-SDK contract, so codec-framed payloads interoperate across SDKs.
taskito/codecs.py:PayloadCodecprotocol,GzipCodec(zip-bomb-guarded decode, 64 MiB default cap),AesGcmCodec([12B nonce][ciphertext||16B tag],encryptionextra),HmacCodec([32B mac][body], constant-time compare), andCodecSerializer(encode in list order, decode in reverse).Queue(codec=...)global chain wraps the queue serializer, covering payloads and results;Queue(codecs={...})+@queue.task(codecs=[...])applies named codecs per task (payload only — results stay on the queue serializer).CryptoError(SerializationError)for decrypt/verify failures.idempotent=Truestays deterministic under per-task encryption codecs.Result-path fix
Worker paths previously wrote results as raw cloudpickle regardless of the configured serializer, while
JobResult.result()read them back with the queue serializer — silently broken for any non-default serializer, and incompatible with a codec chain. All four result-write sites (thread-pool worker and native-async executor in Rust, async_support executor, prefork child) now route through a newQueue._serialize_result(), and the saga orchestrator's result read is aligned to the queue-level serializer.Backward compatibility: results persisted before this change are raw cloudpickle; the default
SmartSerializerreads untagged legacy bytes, so existing queues upgrade cleanly. Non-default serializers could never read results correctly before, so this is strictly an improvement. Jobs persisted before enabling a codec chain cannot be decoded through it — drain the queue first.Tests
tests/core/test_codecs.py): unit round-trips, tamper/short/cap rejection, chain reversibility, worker e2e for global chain and per-task codecs, idempotency-under-encryption regression.tests/core/test_result_serialization.py) plus updated native-async executor tests.